import asyncio
from llama_index.core.agent import ReActAgent
from llama_index.core.tools import FunctionTool
from llama_index.llms.dashscope import DashScope



from llama_index.core.agent.workflow import ReActAgent
from llama_index.llms.deepseek import DeepSeek
import os
from dotenv import load_dotenv


# 加载 API 配置
load_dotenv()
# api_key = os.getenv("ZHIPU_API_KEY")
# api_base_url = os.getenv("ZHIPU_BASE_URL")

api_key = os.getenv("DS_API_KEY")
api_base_url = os.getenv("DS_BASE_URL")

# 选择模型
model = "deepseek-chat"



# 定义一个简单的计算器工具
def multiply(a: float, b: float) -> float:
    """
    简单计算器
    Args:
        a: 第一个数字
        b: 第二个数字
    Returns:
        计算结果
    :param a:
    :param b:
    :return:
    """
    return a * b


# 将函数封装成 LlamaIndex 的工具
# 创建一个代理工作流，将计算器工具传入
agent = ReActAgent(
    tools=[multiply],
    llm=DeepSeek(model=model, api_key=api_key, api_base_url=api_base_url),
    verbose=True,
    system_prompt="""你是一个智能助手，可以调用工具来回答用户问题。
请严格按照以下格式作答：

Question: 用户的问题
Thought: 你的推理过程
Action: 要调用的工具名称（如 multiply）
Action Input: 参数（必须是 JSON 格式）

工具将返回 Observation: 结果

当你知道最终答案时，请使用：
Final Answer: 答案内容

注意：结尾必须严格使用 Final Answer: 开头，不能用 Answer:。
"""
)


async def main():
    # 运行代理
    from llama_index.core.agent.workflow import AgentStream, ToolCallResult

    handler = agent.run("计算一下：20+(2*4)?", )

    async for ev in handler.stream_events():
        if isinstance(ev, AgentStream):
            print(f"{ev.delta}", end="", flush=True)

    response = await handler
    print(response)


# 运行代理
if __name__ == "__main__":
    asyncio.run(main())